#!/usr/bin/env python
# encoding=utf-8
#codeby     道长且阻
#email      ydhcui@suliu.net/QQ664284092
#website    http://www.suliu.net


import time

import requests

#import dns.resolver
from lib.nmap import PortScanner
from core.websearch import BaiduEngine,SogouEngine,SoEngine,BingEngine
from core.base import BaseScan,Request
from core.crawler import Crawler
from core.log import LOG
import models

SUBNAMES = './data/subnames.txt'

class UpdatePortInfo(BaseScan):
    def __init__(self, query):
        self.query = query

    def getver(self,url):
        if not url.startswith(('http','HTTP')):
            url = 'http://'+url
        url = '/'.join(url.strip().split('/')[:3]) #取主机
        try:
            req = requests.get(url+'/SCANOL404TEST.js',timeout=10)
        except Exception as e:
            LOG.error(e)
        headers = req.headers
        text = req.text
        server = headers.get('Server') if 'Server' in headers else headers.get('X-Powered-By')
        if 'Tomcat' in text:
            server = re.findall("<h3>(.*?)</h3>",text)[0]
        if 'Hypertext' in text:
            server = 'Weblogic '+re.findall("<H4>(.*?)404 Not Found</H4>",text)[0]
        return server

    def save(self):
        for Q in self.query:
            url = '%s://%s:%s'%(Q.name,Q.host,Q.port)
            try:
                req = resuests.get()
                HR,created = mod.get_or_create(host=Q.host,port=Q.port)
                HR.createdate = int(time.time())
                HR.headers = req.headers
                HR.content = req.content
                HR.state = req.status_code
                banner = req.headers.get('Server')
                HR.banner = banner if banner else self.getver(url)
                HR.xpoweredby = req.headers.get('X-Powered-By')
                HR.save() #怕获取title出错 所以保存两次
                HR.title = ''.join(re.findall(r"<title>(.*?)</title>",req.text))
                HR.save()
            except Exception as e:
                LOG.error(e)

class PortScan(BaseScan):
    ps = PortScanner()
    def __init__(self,taskid,host,**kwargs):
        host = host.strip()
        args = kwargs.get('args')
        if args:
            self.result = self.ps.scan(host,arguments=args)
        else:
            self.result = self.ps.scan(host)
        M.models.ScanTask
        M..get(M.task_id==taskid)
    def save(self):
        for host,value in self.result.get('scan').items():
            for type in value.all_protocols():
                for port in value.get(type):
                    PR,created = mod.get_or_create(host=host,port=port)
                    PR.type = type
                    PR.name = value.get(type).get(port).get('name')
                    PR.state = value.get(type).get(port).get('state')
                    PR.product = value.get(type).get(port).get('product')
                    PR.extrainfo = value.get(type).get(port).get('extrainfo')
                    PR.reason = value.get(type).get(port).get('reason')
                    PR.version = value.get(type).get(port).get('version')
                    PR.cpe = value.get(type).get(port).get('cpe')
                    PR.createdate = int(time.time())
                    PR.save()

class DomainScan(BaseScan):
    def __init__(self,taskid,host,**kwargs):
        self.namelist = open(SUBNAMES).readlines()
        self.resolvers = Resolver()
        self.answers = []
        self.result = set()
        self.brute(host)
        self.search(host)

    def brute(self,target):
        target = target.strip()
        for subdomain in self.namelist:
            subdomain = subdomain.strip()+'.'+target
            answers = self.resolvers.query(subdomain)
            for answer in answers:
                self.result.add((subdomain,answer.address))

    def search(self,target):
        target = target.strip()
        starttime = time.time()
        E = [e for e in [BaiduEngine,SogouEngine,SoEngine,BingEngine]]
        [e.search('site:%s'%target) for e in E]
        while True:
            endtime = time.time()
            if endtime-starttime>=300: #5分钟
                [e.stop() for e in E]
                if et.Queue.empty():
                    continue
                #return
            url = et.Queue.get()
            netloc = urlparse.urlsplit(url).netloc
            if netloc:
                if ':' in netloc:
                    netloc = netloc.split(':')[0]
                answers = self.recv(netloc)
                for answer in answers:
                    self.result.add((netloc,answer.address))

    def save(self):
        for domain,host in self.result:
            DR,created = mod.get_or_create(domain=domain)
            DR.host = host
            DR.createdate = int(time.time())
            DR.save()
